'''
https://blog.csdn.net/weixin_41227420/article/details/135000156
文档总结的三种方式：Stuff，Refine，Map-Reduce
    1.Stuff:这种方法包括从文档中收集所有信息，并以逻辑简洁的方式进行组织。这包括确定要点、关键思想和支持细节。目标是在不丢失基本信息的情况下，将文档压缩成一个更短、更易于管理的摘要。
    2.Refine:这种方法侧重于通过删除不必要或多余的信息来细化和改进文档摘要。它包括对句子进行编辑和改写，使其更加简洁明了。其目的是创建一个摘要，准确地捕捉文档的主要思想，同时消除任何不必要的细节。
    3.Map-Reduce:这种方法包括将文档分解成更小的部分或块，并将它们分配给不同的个人或团队进行摘要。每个人或小组都有责任总结他们分配的部分。一旦对所有章节进行了总结，这些总结就会合并或缩减为一个连贯的文档总结。这种方法允许以更高效和协作的方式汇总大型复杂文档。

第一种stuff模式总结文章对于短篇文章来说适用，但是篇幅过长就需要很长一段时间。所以今天我们来研究第二种Refine模式来总结文档。

下面我将使用LangChain表达式语言来实现RefineDocumentsChain，使其具有所有内置LCEL功能，其中我们着重点在于流式输出
'''
from functools import partial
from operator import itemgetter

from langchain.callbacks.manager import trace_as_chain_group
# from langchain.chat_models import ChatAnthropic
from langchain_community.llms import Ollama
from langchain.prompts import PromptTemplate
from langchain.schema import StrOutputParser
from langchain_core.prompts import format_document
import os

# 引入LangSmith
os.environ['LANGCHAIN_TRACING_V2'] = 'true'
os.environ['LANGCHAIN_ENDPOINT'] = 'https://api.smith.langchain.com'
os.environ['LANGCHAIN_API_KEY'] = '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'
os.environ['LANGCHAIN_PROJECT'] = 'demo01'

# 导入Ollama的Qwen0.5B的大模型
llm = Ollama(model="qwen:0.5b")

first_prompt = PromptTemplate.from_template("Summarize this content:\n\n{context}")
document_prompt = PromptTemplate.from_template("{page_content}")
partial_format_doc = partial(format_document, prompt=document_prompt)
summary_chain = {"context": partial_format_doc} | first_prompt | llm | StrOutputParser()

refine_prompt = PromptTemplate.from_template(
    "Here's your first summary: {prev_response}. "
    "Now add to it based on the following context: {context}"
)
refine_chain = (
        {
            "prev_response": itemgetter("prev_response"),
            "context": lambda x: partial_format_doc(x["doc"]),
        }
        | refine_prompt
        | llm
        | StrOutputParser()
)


def refine_loop(docs):
    with trace_as_chain_group("refine loop", inputs={"input": docs}) as manager:
        summary = summary_chain.invoke(
            docs[0], config={"callbacks": manager, "run_name": "initial summary"}
        )
        for i, doc in enumerate(docs[1:]):
            summary = refine_chain.invoke(
                {"prev_response": summary, "doc": doc},
                config={"callbacks": manager, "run_name": f"refine {i}"},
            )
        manager.on_chain_end({"output": summary})
    return summary


from langchain.schema import Document

text = """《中国共产党纪律处分条例》第八十二条、第八十三条对搞非组织活动行为列出负面清单，对违规组织、参加乡友会，以及搞拉票、助选、干扰选举活动作出处分规定。
　　1.违规组织、参加乡友会

　　《条例》第八十二条规定，党员领导干部违反有关规定组织、参加自发成立的老乡会、校友会、战友会等，情节严重的，给予相应处分。

　　这里的“违反有关规定”，主要是指违反2002年中央纪委、中央组织部、解放军总政治部印发的《关于领导干部不得参加自发成立的“老乡会”、“校友会”、“战友会”组织的通知》等规定。据此，领导干部不得参加自发成立（未经民政部门登记注册）的老乡、校友、战友之间的各种联谊会之类的组织，不得担当这类联谊会的发起人和组织者，不得在这类联谊会中担任相应职务；领导干部参加老乡、校友、战友之间的各种联谊类活动，不得借机编织“关系网”，搞亲亲疏疏，团团伙伙，更不得有“结盟”“金兰结义”等行为。同时，明确党员领导干部违反上述规定的，党组织要及时进行批评教育，促其改正；对违背组织原则，情节严重的，要给予党纪处分。

　　党员领导干部参加正常范围内的老乡、校友、战友聚会，进行健康文明、积极向上的人际交往，并不违犯党的纪律。但是，一些党员领导干部热衷于参加自发成立的老乡会、校友会、战友会等组织，编织“关系网”，拉“小圈子”，搞非组织活动，在党员群众中产生不良影响，必须对这类行为予以纠正。



　　2.搞拉票、助选、干扰选举活动

　　《条例》第八十三条规定，有下列行为之一的，给予相应处分。

　　①在民主推荐、民主测评、组织考察和党内选举中搞拉票、助选等非组织活动。

　　②在法律规定的投票、选举活动中违背组织原则搞非组织活动，组织、怂恿、诱使他人投票表决。

　　③在选举中进行其他违反党章、其他党内法规和有关章程活动。

　　党内选举是党内政治生活的重要组成部分。《关于新形势下党内政治生活的若干准则》规定，党内选举必须体现选举人意志，坚决防止和查处拉票贿选等行为。在选举活动中搞非组织活动，干扰选举人表达选举意愿，侵犯其他被选举人权益，侵蚀党内政治文化，在党员群众中造成不良影响，对这类违反组织纪律的行为必须予以坚决惩治。

　　搞有组织的拉票贿选，或者用公款拉票贿选的，从重或者加重处分。

　　相较于前述拉票、助选、干扰选举等活动，搞有组织的拉票贿选或者用公款拉票贿选行为，对党的纪律和党员权利的破坏更为严重，影响更为恶劣，因此对此类行为应当予以从重或者加重处分。

　　“七个有之”问题，其中之一就是搞收买人心、拉动选票。拉票贿选直接冲击中国特色社会主义政治制度，动摇党的干部工作根基，危害党的政治生态，损害党和政府形象。《关于新形势下党内政治生活的若干准则》明确规定，“党的任何组织和个人不得以任何方式妨碍选举人依照规定自主行使选举权，坚决反对和防止侵犯党员选举权和被选举权的现象，坚决防止和查处拉票贿选等行为”，“坚决禁止跑官要官、买官卖官、拉票贿选等行为”。
"""

docs = [
    Document(
        page_content=split,
        metadata={"source": "https://en.wikipedia.org/wiki/Nuclear_power_in_space"},
    )
    for split in text.split("\n\n")
]
# 最后调用方法：可以在控制台看见打印输出
print(refine_loop(docs))

# LCEL提供了stream方法和astream方法
# for chunk in runnable.ainvoke({"prev_response": summary, "doc": input_})
# print(chunk,flush=true)